/*
* Copyright (C) 2012-2014 qar
* License: http://www.gnu.org/licenses/gpl.html GPL version 2
*/

#include "crypto/encrypteddevice.h"

#include <QtCore>

#include "outputstream.h"

EncryptedDevice::EncryptedDevice(QIODevice* rawDevice, QObject* parent) :
    QIODevice(parent),
    rawDevice(rawDevice),
    version(maxVersion()),
    clearTextPos(0),
    reachedEnd(false),
    integrityHasher(Hasher::NoIntegrity),
    cipher(Cipher::NoEncryption, QString()),
    authenticator(Authenticator::NoAuthentication, QString())
{}

EncryptedDevice::~EncryptedDevice()
{
    if (isOpen()) {
        qWarn() << "Call EncryptedDevice.close() before the destructor!";
        close();
    }
}

bool EncryptedDevice::open(OpenMode mode,
                           Hasher::IntegrityAlgorithm integrityAlgorithm,
                           Cipher::EncryptionAlgorithm encryptionAlgorithm,
                           Authenticator::AuthenticationAlgorithm authenticationAlgorithm,
                           QString keyFile)
{
    clearTextPos = 0;
    reachedEnd = false;
    if ((mode & QIODevice::Append) ||
        (mode & QIODevice::Text) ||
        (!(mode & QIODevice::ReadWrite)) ||
        ((mode & QIODevice::ReadOnly) && (mode & QIODevice::WriteOnly))) {
        qWarn() << "Incorrect open flags!";
        exit(-1);
        return false;
    }

    if (!rawDevice->isOpen()) {
        if (!rawDevice->open(mode)) {
            qWarn() << "Could not open file!";
            exit(-1);
            return false;
        }
    } else {
        // The read/write flags of the two devices must be identical.
        // The file must not be open in text mode.
        OpenMode mask = QIODevice::ReadWrite | QIODevice::Text | QIODevice::Append;
        if ((rawDevice->openMode() & mask) != (mode & mask)) {
            qWarn() << "Mismatch between open flags!";
            exit(-1);
            return false;
        }
    }

    setOpenMode(mode | QIODevice::Unbuffered);

    if (mode & QIODevice::WriteOnly) {
        // If write, get enc mode as param and write it
        QDataStream out(rawDevice);

        version = maxVersion();
        out << version;

        integrityHasher.setAlgorithm(integrityAlgorithm);
        cipher.setAlgorithm(encryptionAlgorithm, keyFile, true);
        authenticator.setAlgorithm(authenticationAlgorithm, keyFile);

        if (!integrityHasher.toStream(out)) {
            qWarn() << "Could not encode integrityHasher!";
            exit(-1);
            return false;
        }
        if (!cipher.toStream(out)) {
            qWarn() << "Could not encode cipher!";
            exit(-1);
            return false;
        }
        if (!authenticator.toStream(out)) {
            qWarn() << "Could not encode authenticator!";
            exit(-1);
            return false;
        }

        while (rawDevice->pos() % rawBlockSize()) {
            out << quint8(0);
        }
    } else if (mode & QIODevice::ReadOnly) {
        // If read, then read enc mode etc
        QDataStream in(rawDevice);

        in >> version;
        if (version < minVersion() || version > maxVersion()) {
            qWarn() << "Could not decode version number!";
            exit(-1);
            return false;
        }

        if (!integrityHasher.fromStream(in)) {
            qWarn() << "Could not decode hasher!";
            exit(-1);
            return false;
        }

        cipher.setAlgorithm(Cipher::NoEncryption, keyFile, false);
        if (!cipher.fromStream(in)) {
            qWarn() << "Could not decode cipher!";
            exit(-1);
            return false;
        }

        if (!authenticator.fromStream(in)) {
            qWarn() << "Could not decode authenticator!";
            exit(-1);
            return false;
        }

        quint8 v;
        while (rawDevice->pos() % rawBlockSize()) {
            in >> v;
        }
    } else {
        qWarn() << "Incorrect open flags!";
        exit(-1);
        return false;
    }

    qErr() << "Storage version:" << version;
    qErr() << "Storage block size:" << rawBlockSize();
    qErr() << "Storage integrity check:" << integrityHasher.getAlgorithmName();
    qErr() << "Storage encryption:" << cipher.getAlgorithmName();
    qErr() << "Storage authentication:" << authenticator.getAlgorithmName();

    return true;
}

void EncryptedDevice::close()
{
    if (openMode() & QIODevice::WriteOnly) {
        if (!clearBuffer.isEmpty()) {
            writeSegment(clearBuffer);
            clearBuffer.clear();
        }
        writeSegment(QByteArray());
    }
    setOpenMode(NotOpen);
}

bool EncryptedDevice::isSequential() const
{
    return true;
}

qint64 EncryptedDevice::pos() const
{
    return clearTextPos;
}

qint64 EncryptedDevice::size() const
{
    return bytesAvailable();
}

bool EncryptedDevice::seek(qint64)
{
    qWarn() << "Could not seek! Not supported.";
    return false;
}

bool EncryptedDevice::atEnd() const
{
    if (!clearBuffer.isEmpty())
        return false;
    return reachedEnd;
}

bool EncryptedDevice::reset()
{
    qWarn() << "Could not reset stream! Not supported.";
    return false;
}

qint64 EncryptedDevice::bytesAvailable() const
{
    return clearBuffer.size();
}

qint64 EncryptedDevice::bytesToWrite() const
{
    return 0;
}

qint64 EncryptedDevice::writeData(const char* data, qint64 size)
{
    clearBuffer.append(data, size);
    clearTextPos += size;
    while (clearBuffer.size() >= maxClearSegmentSize()) {
        QByteArray segment = clearBuffer.mid(0, maxClearSegmentSize());
        clearBuffer = clearBuffer.mid(segment.size());
        writeSegment(segment);
    }
    return size;
}

qint64 EncryptedDevice::readData(char* data, qint64 size)
{
    while (!atEnd() && clearBuffer.size() < size) {
        if (!readSegment())
            break;
    }

    size_t countRead = qMin(size, (qint64)clearBuffer.size());
    memcpy(data, clearBuffer.constData(), countRead);
    clearBuffer = clearBuffer.mid(countRead);
    clearTextPos += countRead;

    if (clearBuffer.isEmpty() && !atEnd())
        readSegment();

    if (countRead == 0 && atEnd())
        return -1;

    return countRead;
}

bool EncryptedDevice::writeSegment(QByteArray clearSegment)
{
    const bool debug = false;
    if (debug)
        qDebug() << "clearSegment" << clearSegment.size();

    // [cleartext] =>
    // [len(cleartext)][cleartext][padding]
    QByteArray paddedSegment;
    {
        QDataStream out(&paddedSegment, QIODevice::WriteOnly);
        out.setVersion(QDataStream::Qt_4_0);
        out << clearSegment;
        while ((paddedSegment.size() + integrityHashSize()) % encryptionBlockSize() != 0) {
            out << (qint8)(42);
        }
    }

    if (debug)
        qDebug() << "paddedSegment" << paddedSegment.size() << paddedSegment.toHex();

    // [len(cleartext)][cleartext][padding] =>
    // [len(cleartext)][cleartext][padding][hash] : length multiple of encryptionBlockSize()
    // with [hash] = hash([len(cleartext)][cleartext][padding])
    QByteArray hashedPaddedSegment = paddedSegment;
    hashedPaddedSegment.append(integrityHasher.offlineHash(paddedSegment));

    if (debug)
        qDebug() << "hashedPaddedSegment" << hashedPaddedSegment.size() << hashedPaddedSegment.toHex();

    // [len(cleartext)][cleartext][padding][hash] =>
    // [ciphertext                              ]
    QByteArray encryptedHashedPaddedSegment;
    cipher.encrypt(hashedPaddedSegment, encryptedHashedPaddedSegment);

    if (debug)
        qDebug() << "encryptedHashedPaddedSegment" << encryptedHashedPaddedSegment.size() << encryptedHashedPaddedSegment.toHex();

    // [ciphertext                              ] =>
    // [len(ciphertext)][ciphertext                              ][auth]
    // with [auth] = auth([ciphertext                              ])
    QByteArray rawSegment;

    QDataStream out(&rawSegment, QIODevice::WriteOnly);
    out.setVersion(QDataStream::Qt_4_0);
    out << quint32(encryptedHashedPaddedSegment.size());

    if (debug)
        qDebug() << "rawSegment" << rawSegment.size() << rawSegment.toHex();

    QByteArray signedSegment;
    authenticator.signSegment(encryptedHashedPaddedSegment, signedSegment);
    rawSegment.append(signedSegment);

    if (debug)
        qDebug() << "rawSegment" << rawSegment.size() << rawSegment.toHex();

    // [len(ciphertext)][ciphertext                              ][auth] =>
    // [len(ciphertext)][ciphertext                              ][auth][zeros] : length = rawBlockSize()
    while (rawSegment.size() < rawBlockSize()) {
        rawSegment.append(char(0));
    }

    if (debug)
        qDebug() << "rawSegment" << rawSegment.size() << rawSegment.toHex();

    if (writeRawData(rawSegment) != rawSegment.size()) {
        qWarn() << "Could not write segment!";
        return false;
    }

    return true;
}

bool EncryptedDevice::readSegment()
{
    const bool debug = false;
    // rawSegment:
    // [len(ciphertext)][ciphertext                              ][auth][zeros] : length = rawBlockSize()
    QByteArray rawSegment = readRawData(rawBlockSize());

    if (debug)
        qDebug() << "rawSegment" << rawSegment.size() << rawSegment.toHex();

    QDataStream in(&rawSegment, QIODevice::ReadOnly);
    in.setVersion(QDataStream::Qt_4_0);

    // encryptedSegmentSize = len(ciphertext)
    quint32 encryptedSegmentSize;
    in >> encryptedSegmentSize;

    // signedSegment:
    // [ciphertext                              ][auth]
    QByteArray signedSegment;
    if (encryptedSegmentSize + authenticationHashSize() > 0) {
        if (encryptedSegmentSize + authenticationHashSize() <= qint64(rawSegment.size()) - qint64(sizeof(encryptedSegmentSize))) {
            signedSegment.resize(encryptedSegmentSize + authenticationHashSize());
            in.readRawData(signedSegment.data(), encryptedSegmentSize + authenticationHashSize());
        } else {
            qWarn() << "Could not read segment! Size mismatch.";
        }
    }
    if (in.status() != QDataStream::Ok) {
        qWarn() << "Could not read segment!";
        return false;
    }

    if (debug)
        qDebug() << "signedSegment" << signedSegment.size() << signedSegment.toHex();

    // encryptedHashedPaddedSegment:
    // [ciphertext                              ]
    QByteArray encryptedHashedPaddedSegment;
    if (!authenticator.checkAndExtractSegment(signedSegment, encryptedHashedPaddedSegment)) {
        qWarn() << "Authenticity check for segment failed!";
        return false;
    }

    if (debug)
        qDebug() << "encryptedHashedPaddedSegment" << encryptedHashedPaddedSegment.size() << encryptedHashedPaddedSegment.toHex();

    // hashedPaddedSegment:
    // [len(cleartext)][cleartext][padding][hash]
    QByteArray hashedPaddedSegment;
    cipher.decrypt(encryptedHashedPaddedSegment, hashedPaddedSegment);

    if (debug)
        qDebug() << "hashedPaddedSegment" << hashedPaddedSegment.size() << hashedPaddedSegment.toHex();

    if (hashedPaddedSegment.size() < integrityHashSize()) {
        qWarn() << "Could not read integrity hash value!";
        return false;
    }

    // paddedSegment:
    // [len(cleartext)][cleartext][padding]
    QByteArray paddedSegment = hashedPaddedSegment.left(hashedPaddedSegment.size() - integrityHashSize());

    if (debug)
        qDebug() << "paddedSegment" << paddedSegment.size() << paddedSegment.toHex();

    // integrityHashRead:
    // [hash]
    QByteArray integrityHashRead = hashedPaddedSegment.right(integrityHashSize());

    // integrityHashComputed = hash([len(cleartext)][cleartext][padding])
    QByteArray integrityHashComputed = integrityHasher.offlineHash(paddedSegment);
    if (integrityHashRead != integrityHashComputed) {
        qWarn() << "Integrity hash mismatch!";
        return false;
    }

    // clearSegment:
    // [cleartext]
    QByteArray clearSegment;
    {
        QDataStream in(&paddedSegment, QIODevice::ReadOnly);
        in.setVersion(QDataStream::Qt_4_0);
        in >> clearSegment;
        if (in.status() != QDataStream::Ok) {
            qWarn() << "Read error!";
            return false;
        }
        while (!in.atEnd()) {
            qint8 padding;
            in >> padding;
            if (padding != 42) {
                qWarn() << "Read error!";
                return false;
            }
        }
    }

    if (debug)
        qDebug() << "clearSegment" << clearSegment.size() << clearSegment.toHex();

    if (clearSegment.isEmpty()) {
        reachedEnd = true;
    }

    clearBuffer.append(clearSegment);
    return true;
}

qint64 EncryptedDevice::writeRawData(QByteArray data)
{
    qint64 totalWritten = 0;
    for (int i = 0; i < data.size();) {
        qint64 chunk = rawDevice->write(data.constData() + i, data.size() - i);
        if (chunk <= 0) {
            qWarn() << "Probably a write error!";
            break;
        }
        i += chunk;
        totalWritten += chunk;
    }
    return totalWritten;
}

QByteArray EncryptedDevice::readRawData(int len)
{
    QByteArray data(len, 0);

    int size = 0;
    while (size < len) {
        qint64 chunk = rawDevice->read(data.data() + size, len - size);
        if (chunk <= 0) {
            qWarn() << "Probably a write error!";
            break;
        }
        size += chunk;
    }
    data.resize(size);
    return data;
}

qint64 EncryptedDevice::integrityHashSize() const
{
    return integrityHasher.hashSize();
}

qint64 EncryptedDevice::encryptionBlockSize() const
{
    return cipher.blockSize();
}

qint64 EncryptedDevice::authenticationHashSize() const
{
    return authenticator.hashSize();
}

qint64 EncryptedDevice::maxClearSegmentSize() const
{
    return maxEncryptedSegmentSize() - sizeof(qint32) - integrityHashSize();
}

qint64 EncryptedDevice::maxEncryptedSegmentSize() const
{
    qint64 result = rawBlockSize() - sizeof(quint32) - authenticationHashSize();
    result = (result / encryptionBlockSize()) * encryptionBlockSize();
    while (result > rawBlockSize()) {
        result -= encryptionBlockSize();
    }
    return result;
}

qint64 EncryptedDevice::rawBlockSize() const
{
    if (version >= 1)
        return 16 * 1024;
    return 16 * 1024;
}

qint32 EncryptedDevice::minVersion()
{
    return 1;
}

qint32 EncryptedDevice::maxVersion()
{
    return 1;
}
